import sys
import os
sys.path.append("/home/qiufengfeng/nlp/nlp_project/gitee_my/nlp_tools")
from nlp_tools.corpus import ChineseDailyNerCorpus

from nlp_tools.callbacks.eval_callback import F1ScoreSaveCallBack

from nlp_tools.embeddings.transformer_embedding import TransformerEmbedding
train_x,train_y = ChineseDailyNerCorpus.load_data('train')
test_x,test_y = ChineseDailyNerCorpus.load_data('test')
valid_x,valid_y = ChineseDailyNerCorpus.load_data('valid')

from nlp_tools.tasks.labeling import BiLSTM_CRF_Model

import sys
if 'win' in sys.platform:
    bert_model_path = r"E:\nlp-data\pretrain_model\bert\model_uncase"
    save_path = r'E:\github\output_models\ner.h5'
else:
    bert_model_path = "/home/qiufengfeng/nlp/pre_trained_model/chinese_L-12_H-768_A-12/chinese_L-12_H-768_A-12"
    save_path = '/home/qiufengfeng/nlp/train_models/ner/'

config_path = os.path.join(bert_model_path,'bert_config.json')
check_point_path = os.path.join(bert_model_path,'bert_model.ckpt')
vocab_path = os.path.join(bert_model_path,'vocab.txt')
embedding = TransformerEmbedding(vocab_path=vocab_path,config_path=config_path,checkpoint_path=check_point_path,task='labeling')


model = BiLSTM_CRF_Model(embedding)
model.fit(train_x,train_y,epochs=10,callbacks=[F1ScoreSaveCallBack(model,valid_x,valid_y,step=1,batch_size=128,model_save_path=save_path)])


#model.evaluate(test_x,test_y)

import time
from sklearn.metrics import f1_score
start = time.time()
y_predict = model.predict(test_x)
end = time.time()
print("cost:"+ str(end - start))
print(model.evaluate(test_x,test_y))